# only face images, no target or label
from config import HP
from torchvision import transforms as T  # torch audio
from torch.utils.data import DataLoader
import torchvision.datasets as TD
import os
import matplotlib.pyplot as plt
import torchvision.utils as vutil

os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'  # openKMP cause unexpected error

# apply
data_face = TD.ImageFolder(
    root=HP.data_root,
    transform=T.Compose([
        T.Resize(HP.image_size),
        T.CenterCrop(HP.image_size),
        T.ToTensor(),
        T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
)

face_loader = DataLoader(
    data_face,
    batch_size=HP.batch_size,
    shuffle=True,
    num_workers=HP.n_workers
)


# de-normalize
invTrans = T.Compose([
    T.Normalize(mean=[0., 0., 0.], std=[1/0.5, 1/0.5, 1/0.5]),
    T.Normalize(mean=[-0.5, -0.5, -0.5], std=[1., 1., 1.])
])

if __name__ == '__main__':
    for data, _ in face_loader:
        print(data.size())
        # format into 8 * 8
        grid = vutil.make_grid(data, nrow=12)  # 将小图片设计成大图片
        plt.imshow(invTrans(grid).permute(1, 2, 0))
        plt.show()
        break
